// Copyright (c) 2024 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package lm

import (
	"fmt"

	"github.com/urfave/cli/v2"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

type Labels map[string]string

func (labels Labels) Labels() Labels {
	return labels
}

func (labels Labels) Output(c *cli.Context) error {
	client, node, err := getNode(c)
	if err != nil {
		return err
	}

	nodeCopy := node.DeepCopy()
	updated := false

	var removeLabelKeys []string

	if !checkDevicePlugin(node) {
		removeLabelKeys = append(removeLabelKeys, hostArchLabelKey, acceleratorLabelKey, acceleratorTypeLabelKey)
	}

	if _, exists := labels[npuPresentLabelKey]; !exists {
		removeLabelKeys = append(removeLabelKeys, npuPresentLabelKey, masterLabelKey,
			workerLabelKey, workerRoleLabelKey)
	}

	if _, exists := labels[runtimeLabelKey]; !exists {
		removeLabelKeys = append(removeLabelKeys, runtimeLabelKey)
	}

	updated = removeLabels(nodeCopy.Labels, removeLabelKeys)

	for key, value := range labels {
		updated = addLabel(nodeCopy.Labels, key, value) || updated
	}

	if updated {
		if _, err = client.CoreV1().Nodes().Update(c.Context, nodeCopy, metav1.UpdateOptions{}); err != nil {
			return fmt.Errorf("update node labels: %w", err)
		}
	}
	return nil
}

func addLabel(labels map[string]string, key, value string) bool {
	if labels == nil {
		return false
	}
	if v, ok := labels[key]; ok && v == value {
		return false
	}
	labels[key] = value
	return true
}

func removeLabels(labels map[string]string, removeLabelKeys []string) bool {
	removed := false

	for _, key := range removeLabelKeys {
		if _, ok := labels[key]; ok {
			removed = true
			delete(labels, key)
		}
	}
	return removed
}
